library(Matrix)

args=commandArgs(T)
GCMatFile=args[1] #GCMat(must contain BarXCRFile's barcode!)
BarXCRFile=args[2] #Bar XCR info Table!
XCRflag=args[3]  #1,TCR;2,BCR
GeneListFile=args[4]   #if "F",then use the gene list!
BarClusFile=args[5]  #if "F",then use the barcode only BarXCRFile! File Format:barcode cluster
BarMode=args[6] #1,use XCR Table barcode; 2, use BarCluster File barcode
CITERData=args[7] #if not, fill "F". AllSCITEBarExpMat must in RData

OutGeneExpFile=args[8]

#test:
#setwd("/data2/users/dongxingjun/Human_SC_Ana/TCRdg/B2/P23-2")
#GCMatFile="P23-2.GCMat.RData"
#BarXCRFile="P23-2.BarTCRInfo.xls"
#XCRflag="1"
#GeneListFile="/home/xingjun/database/code/SingleCell1/GeneExpTable/GeneList/GeneList0927.xls"
#BarClusFile = "F"
#BarMode = "1"
#CITERData = "F"
#OutGeneExpFile="P23-2.GeneExp.Test.xls"

#finish in 2020.4.20

#combine Barcode's GCMat,XCR Info Table,Cluster Table,..., create a new table!
#new table:UMI,GeneExpNum,XCRInfo,GeneExpProfile,... and so on
#if TCR's barcode not in CDNA GCMat, it will add in getGCMat.R

#read files
load(GCMatFile)

#BarXCRTable
readTableToDF <- function(TableFile)
{
  
  BarXCRTable <- read.table(BarXCRFile,sep = "\t",header = FALSE,stringsAsFactors = FALSE)
  HeadLine <- as.character(BarXCRTable[1,])
  colnames(BarXCRTable) <- HeadLine
  if(dim(BarXCRTable)[1] == 1)
  {
    BarXCRTable <- BarXCRTable[-1,]
  }else
  {
    BarXCRTable <- BarXCRTable[2:dim(BarXCRTable)[1],]
  }
  
  return(BarXCRTable)
}

BarXCRTable <- readTableToDF(BarXCRFile)
BarXCRTableBarcode <- BarXCRTable$barcode
#BarXCRTable end.

if(length(BarXCRTableBarcode) == 0)
{
  print("ERROR:BarXCRTable don't have barcode!")
  return(FALSE)
}

#geneList
if(GeneListFile != "F")
{
  GeneList1 <- readLines(GeneListFile)
  GCMatGene <- rownames(GCMat)
  GeneList <- intersect(GeneList1,GCMatGene)
}else
{
  GeneList <- c("CD4","CD8A","CD8B","CD79A","CD79B","CD19")
}
#geneList end.

#Barcode - Cluster
if(BarClusFile != "F")
{
  #cluster name add prifix "c"
  BarClusTable <- read.table(BarClusFile,sep = "\t",header = TRUE,stringsAsFactors = FALSE)
  BarClus <- paste("c",BarClusTable$cluster,sep="")
  BarClusTable$cluster <- BarClus
  #add end.
}else
{
  BarClusTable <- matrix("-",ncol = 1,nrow = length(BarXCRTableBarcode))
  BarClusTable <- as.data.frame(cbind(BarXCRTableBarcode,BarClusTable))
  names(BarClusTable) <- c("barcode","cluster")
}

#read files end.

#check barcode (XCRTable barcode must all in GCMat)
GCMatBarcode <- colnames(GCMat)
commonBar <- intersect(BarXCRTableBarcode,GCMatBarcode)
if(length(commonBar) != length(BarXCRTableBarcode))
{
  NoHaveBar <- setdiff(BarXCRTableBarcode,commonBar)
  stri <- paste("WARNING:barcode error! barcode not in GCMat number:",length(NoHaveBar),sep="")
  print(stri)
  Addmat <- matrix(0,nrow=dim(GCMat)[1],ncol=length(NoHaveBar))
  rownames(Addmat) <- rownames(GCMat)
  colnames(Addmat) <- NoHaveBar
  GCMat <- cbind(GCMat,Addmat)
  commonBar <- intersect(BarXCRTableBarcode,GCMatBarcode)
}
#check barcode end!

#base function
AddGeneExp <- function(TempTable,GeneList,GCMat)
{
  TableBar <- as.character(TempTable$barcode)
  ChoMat <- as.matrix(t(GCMat[GeneList,TableBar]))
  UMI <- colSums(GCMat[,TableBar])
  GeneExpNum <- colSums(GCMat[,TableBar] > 0)

  TempTable <- cbind(TempTable,UMI)
  TempTable <- cbind(TempTable,GeneExpNum)
  TempTable <- cbind(TempTable,ChoMat)  

  return(TempTable)
}


TCRSta <- function(BarXCRTable,GCMat,commonBar)
{
#BarXCRTable col:TCR/BCR
    #test input:
    #test input end.
    
    #get temp Gene Mat
    TempGene <- c("CD4","CD8A","CD8B","CD79A","CD79B","CD19")
    ChoMat <- as.matrix(t(GCMat[TempGene,commonBar]))
    NewTable <- cbind(BarXCRTable,ChoMat)
 
    #TCR CD4 or CD8
    AllTCR <- unique(NewTable$TCRName)
    
    TCRChar1List <- c()
    for(i in AllTCR)
    {
      ThisTCRTable <- NewTable[which(NewTable$TCRName == i),]
      CD8ExpNum = sum(ThisTCRTable$CD8A > 0 | ThisTCRTable$CD8B > 0)
      CD4ExpNum = sum(ThisTCRTable$CD4 > 0)
  
      if(CD4ExpNum > CD8ExpNum)
      {
        TCRChar1 = "CD4TCR"
      }else if(CD8ExpNum > CD4ExpNum)
      {
        TCRChar1 = "CD8TCR"
      }else
      {
        TCRChar1 = "unknow"
      }
      TCRChar1List <- c(TCRChar1List,TCRChar1)
    }
    names(TCRChar1List) <- AllTCR
  
    TCRChar1List2 <- c()
    for(i in commonBar)
    {
      ThisBarTCR <- NewTable[which(NewTable$barcode == i),]$TCRName[1]
      TCRChar1List2 <- c(TCRChar1List2,as.character(TCRChar1List[ThisBarTCR]))
    }
    #TCR CD4 or CD8 end.
    
    #TCR tumor special
    TCRChar3List <- c()  #every TCR tumor special info
    
    TableHead <- colnames(NewTable)
    if(c("Tumor") %in% TableHead & c("Normal") %in% TableHead & c("Blood") %in% TableHead)
    {
        tflag <- 1
    }else if(c("Tumor") %in% TableHead & c("Normal") %in% TableHead & !c("Blood") %in% TableHead)
    {
        tflag <- 2
    }else if(c("Tumor") %in% TableHead & !c("Normal") %in% TableHead & c("Blood") %in% TableHead)
    {
        tflag <- 3
    }else
    {
        tflag <- 0
    }

    
    if(tflag > 0)
    {
        for(i in AllTCR)
        {

            ThisTCRTable <- NewTable[which(NewTable$TCRName == i),]
      
            TumorNum <- as.integer(ThisTCRTable$Tumor[1])
            if(tflag == 1)
            {
                ControlNum <- as.integer(ThisTCRTable$Normal[1]) + as.integer(ThisTCRTable$Blood[1])
            }else if(tflag == 2)
            {
                ControlNum <- as.integer(ThisTCRTable$Normal[1])
            }else if(tflag == 3)
            {
                ControlNum <- as.integer(ThisTCRTable$Blood[1])
            }else
            {
                ControlNum <- TumorNum
                print("error in TCR tumor special!")
            }

            if(TumorNum > 5 & ControlNum < 3)
            { 
                TCRChar3 <- "TumorSpecial"
            }else
            {
                TCRChar3 <- "-"
            }

            TCRChar3List <- c(TCRChar3List,TCRChar3)
        }   
 
        #final view 
        names(TCRChar3List) <- AllTCR
        TCRChar3List2 <- c() #final View TCRinfo in Table
        for(i in commonBar)
        {
          ThisBarTCR <- NewTable[which(NewTable$barcode == i),]$TCRName[1]
          TCRChar3List2 <- c(TCRChar3List2,as.character(TCRChar3List[ThisBarTCR]))
        }
    }else
    {
        TCRChar3List2 <- rep("-",time = length(commonBar))
    }
    #TCR tumor special end.

    #TCR doublet
    TCRChar2List <- c()
    for(i in commonBar)
    {
      ThisTCRTable <- NewTable[which(NewTable$barcode == i),]
      CD4Expflag <- 0
      CD8Expflag <- 0
      BMExpflag <- 0
      
      if(ThisTCRTable$CD4 > 0)
      {
        CD4Expflag = 1
      }
      
      if(ThisTCRTable$CD8A > 0 | ThisTCRTable$CD8B > 0)
      {
        CD8Expflag = 1
      }
      
      if(ThisTCRTable$CD79A > 0 | ThisTCRTable$CD79B > 0 | ThisTCRTable$CD19 > 0)
      {
        BMExpflag = 1
      }
      
      if(CD4Expflag == 1 & CD8Expflag == 1)
      {
        TCRChar2 <- "doublet"
      }else if(BMExpflag == 1)
      {
        TCRChar2 <- "doublet"
      }else
      {
        TCRChar2 <- "-"
      }
      TCRChar2List <- c(TCRChar2List,TCRChar2)
    }
    #TCR doublet end.

    #add these col in Table
    RawHead <- colnames(BarXCRTable)
    NewHead <- c(RawHead,"TCRatt1","TCRatt2","TCRatt3")
    NewTable2 <- cbind(BarXCRTable,TCRChar1List2,TCRChar2List,TCRChar3List2)
    colnames(NewTable2) <- NewHead
    #add end.

    return(NewTable2)
}

DataFramAddNull <- function(PreBar,TempTable)
{
  BarTempTable <- TempTable$barcode
  ThisCommBar <- intersect(PreBar,BarTempTable)
  ATable <- BarTempTable[ThisCommBar,]
  RemainBar <- setdiff(PreBar,ThisCommBar)

  if(length(RemainBar) > 0)
  {
    BTable <- as.data.frame(cbind(RemainBar,matrix("-",nrow = RemainBar,ncol=dim(TempTable)[2]-1)))
    FTable <- rbind(ATable,BTable)
  }else
  {
    FTable <- ATable
  }

  return(FTable)
}


FindTableInfo <- function(keyList,Table)
{

  finalTable <- NULL
  FillNULL <- rep("-",time=dim(Table)[2])
  
  TableKey <- Table[,1]
  for(i in keyList)
  {
    if(i %in% TableKey)
    {
      finalTable <- rbind(finalTable,as.matrix(Table[which(Table$barcode == i),][1,]))
    }else
    {
      finalTable <- rbind(finalTable,FillNULL)
    }
    
  }
  
  
  finalTable <- finalTable[,2:dim(finalTable)[2]]
  finalTable <- as.data.frame(finalTable)
  
  return(finalTable)
}

TableMerge1 <- function(XCRTable,ClusTable,flag)
{

  if(flag == "1")
  {
    Prekey <- XCRTable[,1]
    XCRTableHead <- colnames(XCRTable)
    ClusTableHead <- colnames(ClusTable)
    Head <- c("barcode",ClusTableHead[2:length(ClusTableHead)],XCRTableHead[2:length(XCRTableHead)])
    finalTable <- cbind(Prekey,FindTableInfo(Prekey,ClusTable),FindTableInfo(Prekey,XCRTable))
    colnames(finalTable) <- Head
    return(finalTable)
    
  }else if(flag == "2")
  {
    Prekey <- ClusTable[,1]
    XCRTableHead <- colnames(XCRTable)
    ClusTableHead <- colnames(ClusTable)
    Head <- c("barcode",ClusTableHead[2:length(ClusTableHead)],XCRTableHead[2:length(XCRTableHead)])
    finalTable <- cbind(Prekey,FindTableInfo(Prekey,ClusTable),FindTableInfo(Prekey,XCRTable))
    colnames(finalTable) <- Head
    return(finalTable)
    
  }else
  {
    
    return(FALSE)
  }
}
#base function end.

#main

#add TCR att1/att2
if(XCRflag == "1")  #TCR
{
  NewTable2 <- TCRSta(BarXCRTable,GCMat,commonBar)
}else #BCR
{
  NewTable2 <- BarXCRTable
}

#add cluster info 
NewTable3 <- TableMerge1(NewTable2,BarClusTable,BarMode)
  
#add GeneExp
NewTable4 <- AddGeneExp(NewTable3,GeneList,GCMat)

#add CITESEQ
if(CITERData != "F")
{
  load(CITERData)
  Head <- colnames(NewTable4)
  Head <- c(Head,colnames(AllSCITEBarExpMat))
  TableBar <- as.character(NewTable4$barcode)  ###why?
  TableAndCITECommBar <- intersect(TableBar,rownames(AllSCITEBarExpMat))

  NeedAddCITEBar <- setdiff(TableBar,TableAndCITECommBar)
  if(length(NeedAddCITEBar) > 0)
  {
    tempMat <- matrix(0,nrow = length(NeedAddCITEBar),ncol = dim(AllSCITEBarExpMat)[2])
    rownames(tempMat) <- NeedAddCITEBar
    AllSCITEBarExpMat <- rbind(AllSCITEBarExpMat,tempMat)
  }

  NewTable4 <- cbind(NewTable4,AllSCITEBarExpMat[TableBar,])
  colnames(NewTable4) <- Head
}
#combine final Table
write.table(NewTable4,OutGeneExpFile,sep="\t",quote=F,row.names = FALSE,col.names = TRUE)
